import numpy as np
import random
from gym_torcs import TorcsEnv
import argparse
import collections
import math
import pid
# import ipdb


state_size = 29
action_size = 3
train_indicator = 0  # train or not

VISION = False

disController = pid.pidController()
angleController = pid.pidController()

disController.setPID(0.25, 0.05, 2.0)
angleController.setPID(0.45, 0.05, 2.5)

# env environment
env = TorcsEnv(vision=VISION, throttle=True, gear_change=False)


file_reward = open("results/results.txt", "w")
file_distances = open("results/distances.txt", "w")
file_angle = open("results/angle.txt", "w")
file_trackPos = open("results/trackPos.txt", "w")
file_speedX = open("results/speedX.txt", "w")
file_speedY = open("results/speedY.txt", "w")
file_speedZ = open("results/speedZ.txt", "w")
file_track = open("results/track.txt", "w")
file_curve = open("results/curve.txt", "w")
file_action = open("results/action.txt", "w")

steps = 0

for i in range(2000):
    reset_flag = 0

    if np.mod(i, 3) == 0:
        ob, distFromStart = env.reset(relaunch=True)
    else:
        ob, distFromStart = env.reset()

    disController.clear()
    angleController.clear()

    s_t = np.hstack((ob.angle, ob.track, ob.trackPos, ob.speedX,
                    ob.speedY, ob.speedZ, ob.wheelSpinVel / 100.0, ob.rpm))
    for j in range(200000):

        a_t = np.zeros([1, action_size])
        noise_t = np.zeros([1, action_size])
        # ipdb.set_trace()

        distErr = 0.0 - ob.trackPos
        angleShift = disController.nextOutput(distErr)

        angleErr = 0.0 + angleShift - ob.angle

        a_t[0][0] = angleController.nextOutput(angleErr)

        if a_t[0][0] > 1:
            a_t[0][0] = 1
        if a_t[0][0] < (-1):
            a_t[0][0] = -1
        if ob.speedX > 0.2:
            a_t[0][1] = 0
        else:
            a_t[0][1] = 0.5
        a_t[0][2] = 0

        ob, distFromStart, r_t, r_t0, done, info = env.step(a_t[0])

        s_t1 = np.hstack(
            (ob.angle, ob.track, ob.trackPos, ob.speedX, ob.speedY, ob.speedZ, ob.wheelSpinVel / 100.0, ob.rpm))

        s_t = s_t1
        steps += 1
        # file_reward.write(
        #     str(i) + " " + str(steps) + " " + str(distFromStart) + " " + str(r_t) + " " + str(r_t0) + " \n")
        # file_angle.write(str(i) + " " + str(steps) + " " + str(ob.angle) + "\n")
        # file_trackPos.write(str(i) + " " + str(steps) + " " + str(ob.trackPos) + "\n")
        # file_track.write(str(i) + " " + str(steps) + " " + str(ob.track) + "\n")
        # file_speedX.write(str(i) + " " + str(steps) + " " + str(ob.speedX) + "\n")
        # file_speedY.write(str(i) + " " + str(steps) + " " + str(ob.speedY) + "\n")
        # file_speedZ.write(str(i) + " " + str(steps) + " " + str(ob.speedZ) + "\n")
        # file_action.write(
        #     str(i) + " " + str(steps) + " " + str(a_t[0][0]) + " " + str(a_t[0][1]) + " " + str(a_t[0][2]) + "\n")
        if done:
            print("I'm dead")
            break
    file_distances.write(str(i) + " " + str(distFromStart) + "\n")

file_reward.close()
file_distances.close()
file_angle.close()
file_trackPos.close()
file_speedX.close()
file_speedY.close()
file_speedZ.close()
file_track.close()
file_curve.close()
file_action.close()
env.end()
print("Finish.")

# for param in critic.parameters(): param.grad.data.clamp(-1, 1)
